import psycopg2
from psycopg2 import pool
import logging
from psycopg2.extras import RealDictCursor
from contextlib import contextmanager

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 数据库连接配置
DB_HOST = '192.168.186.126'
DB_PORT = '5432'
DB_USER = 'cywq'
DB_PASSWORD = '1234567890'
DB_NAME = 'xaunke_database'

# 数据库连接池
db_pool = None

def init_db_pool():
    """初始化数据库连接池"""
    global db_pool
    try:
        logger.info("正在初始化数据库连接池...")
        logger.info(f"尝试连接到数据库: host={DB_HOST}, port={DB_PORT}, database={DB_NAME}, user={DB_USER}")
        db_pool = pool.SimpleConnectionPool(
            minconn=1,
            maxconn=20,  # 增加最大连接数
            host=DB_HOST,
            port=int(DB_PORT),  # 确保端口是整数
            database=DB_NAME,
            user=DB_USER,
            password=DB_PASSWORD
        )
        # 测试连接
        with get_db_connection() as conn:
            with conn.cursor() as cursor:
                cursor.execute('SELECT 1')
                logger.info("数据库连接测试成功")
        logger.info("数据库连接池初始化成功")
    except Exception as e:
        logger.error(f"数据库连接池初始化失败: {str(e)}")
        raise e

@contextmanager
def get_db_connection():
    """获取数据库连接的上下文管理器"""
    conn = None
    try:
        if db_pool is None:
            init_db_pool()
        conn = db_pool.getconn()
        yield conn
        conn.commit()
    except Exception as e:
        if conn:
            conn.rollback()
        raise e
    finally:
        if conn:
            db_pool.putconn(conn)

@contextmanager
def get_db_cursor():
    """获取数据库游标的上下文管理器"""
    with get_db_connection() as conn:
        with conn.cursor() as cursor:
            yield cursor

def get_db():
    """获取数据库连接（兼容旧代码）"""
    if db_pool is None:
        init_db_pool()
    return db_pool.getconn()

def close_db(conn):
    """关闭数据库连接（兼容旧代码）"""
    if conn is not None:
        db_pool.putconn(conn)

def close_all_connections():
    """关闭所有数据库连接"""
    if db_pool is not None:
        db_pool.closeall()
        logger.info("所有数据库连接已关闭") 